#!/bin/bash
set -e

# 利用msopgen生成可编译文件
rm -rf ./flash_attention
msopgen gen -i flash_attention.json -f tf -c ai_core-Ascend910B2 -lan cpp -out ./flash_attention -m 0 -op FlashAttentionScore
rm -rf flash_attention/framework/tf_plugin
cp -rf tf_plugin ./flash_attention/framework/

cd flash_attention

# 判断当前目录下是否存在CMakePresets.json文件
if [ ! -f "CMakePresets.json" ]; then
  echo "ERROR, CMakePresets.json file not exist."
  exit 1
fi

# 禁止生成CRC校验和
sed -i 's/--nomd5/--nomd5 --nocrc/g' ./cmake/makeself.cmake

# 修改cann安装路径
sed -i 's:"/usr/local/Ascend/latest":"/usr/local/Ascend/ascend-toolkit/latest":g' CMakePresets.json
# 修改vendor_name 防止覆盖之前vendor_name为customize的算子;
# vendor_name需要和aclnn中的CMakeLists.txt中的CUST_PKG_PATH值同步，不同步aclnn会调用失败;
# vendor_name字段值不能包含customize；包含会导致多算子部署场景CANN的vendors路径下config.ini文件内容截取错误
sed -i 's:"customize":"flash_attention":g' CMakePresets.json

bash build.sh

# # 安装编译成功的算子包
bash ./build_out/custom_opp*.run
cd ..
rm -rf ./flash_attention
echo "======== Build & Install customize op success ========="

echo "======== Start to build tf_ops ========="
if [ -d build ];then
    rm -rf build
fi

cmake -S . -B build -G "Unix Makefiles" -DTF_VERSION=$TF_VERSION -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXPORT_COMPILE_COMMANDS=on
cmake --build build -j 100
echo "======== Build tf_ops success ========="

echo "======== Start to remove empty op impl ========="
CUST_OP_PATH=$(find /usr/local/Ascend -name flash_attention)
echo "customize op install path: $CUST_OP_PATH"
cd $CUST_OP_PATH
rm -rf op*
echo "======== Remove empty op impl success ========="